import pandas as pd
from datetime import datetime
from typing import Dict, Type
from ..core.base_strategy import BaseStrategy, BacktestResult
from concurrent.futures import ThreadPoolExecutor
import asyncio
from dask.distributed import Client

class BacktestEngine:
    """分布式回测引擎"""
    
    def __init__(self, data_service, factor_service, dask_cluster='127.0.0.1:8786'):
        self.data_service = data_service
        self.factor_service = factor_service
        self.dask_client = Client(dask_cluster)
        
    async def run_backtest(
        self,
        strategy_class: Type[BaseStrategy],
        params: Dict,
        symbols: List[str],
        start: datetime,
        end: datetime,
        initial_capital: float = 1e6
    ) -> BacktestResult:
        """执行回测"""
        # 获取历史数据
        hist_data = await self.data_service.get_history(
            symbols, start, end
        )
        
        # 获取所需因子
        required_factors = strategy_class.required_factors()
        factor_data = await self._load_factor_data(required_factors, symbols, start, end)
        
        # 初始化策略
        strategy = strategy_class(params)
        strategy.initialize(initial_capital)
        
        # 事件循环
        result = BacktestResult()
        dates = pd.date_range(start, end)
        
        # 使用多线程处理每日数据
        date_chunks = np.array_split(dates, len(self.dask_client.ncores())*4)
        futures = [self.dask_client.submit(self._process_date_chunk, chunk, hist_data, factor_data, strategy) 
                  for chunk in date_chunks]
        
        results = await asyncio.gather(*[future.result() for future in futures])
        return self._calculate_metrics(result)
        
    async def _load_factor_data(self, factors: List[str], symbols: List[str], 
                              start: datetime, end: datetime) -> Dict:
        """加载因子数据"""
        factor_dict = {}
        for factor in factors:
            data = await self.factor_service.get_factor(
                factor, symbols, start, end
            )
            factor_dict[factor] = data
        return factor_dict

    def _process_date_chunk(self, dates_chunk, hist_data, factor_data, strategy):
        """处理日期块"""
        chunk_result = BacktestResult()
        for dt in dates_chunk:
            self._process_day(dt, hist_data, factor_data, strategy, chunk_result)
        return chunk_result

    # 保留原有_process_day和_execute_trades方法
    
    def _process_day(self, dt, hist_data, factor_data, strategy, result):
        """处理单日回测逻辑"""
        daily_data = hist_data.xs(dt, level='date')
        factors = {
            factor: factor_data[factor].xs(dt, level='date')
            for factor in factor_data
        }
        
        # 生成信号
        signals = strategy.generate_signal(daily_data, factors)
        
        # 执行交易
        trades = self._execute_trades(strategy, signals, dt)
        
        # 更新结果
        result.trades.extend(trades)
        result.positions[dt] = strategy.position.copy()
        result.returns.append(strategy.portfolio_value / strategy.initial_capital - 1)
        
    def _execute_trades(self, strategy, signals, dt):
        """模拟交易执行"""
        trades = []
        for _, row in signals.iterrows():
            symbol = row['symbol']
            signal = row['signal']
            price = row['price']
            
            # 简单执行逻辑，实际应包含滑点和手续费
            if signal > 0:  # 买入信号
                shares = int(signal * strategy.cash / price)
                strategy.position[symbol] = strategy.position.get(symbol, 0) + shares
                strategy.cash -= shares * price
                trades.append({
                    'date': dt,
                    'symbol': symbol,
                    'direction': 'BUY',
                    'shares': shares,
                    'price': price
                })
            elif signal < 0:  # 卖出信号
                shares = min(int(-signal), strategy.position.get(symbol, 0))
                if shares > 0:
                    strategy.position[symbol] -= shares
                    strategy.cash += shares * price
                    trades.append({
                        'date': dt,
                        'symbol': symbol,
                        'direction': 'SELL',
                        'shares': shares,
                        'price': price
                    })
        return trades
    
    def _calculate_metrics(self, result: BacktestResult) -> BacktestResult:
        """计算绩效指标"""
        returns = pd.Series(result.returns)
        result.metrics['total_return'] = returns.iloc[-1] - returns.iloc[0]
        result.metrics['annual_return'] = returns.mean() * 252
        result.metrics['sharpe'] = returns.mean() / returns.std() * np.sqrt(252)
        result.metrics['max_drawdown'] = (returns.expanding().max() - returns).max()
        return result